package GreedyAlgorithm;

import java.util.Comparator;
import java.util.PriorityQueue;

/**
 * 使用堆结构，快速的获取中位数
 *
 * @author Liaorun
 */
public class GetMedian {


    public PriorityQueue<Integer> minPartQ = new PriorityQueue<>();
    public PriorityQueue<Integer> maxPartQ = new PriorityQueue<>(new MaxPartComparator());


    public class MaxPartComparator implements Comparator<Integer> {

        @Override
        public int compare(Integer o1, Integer o2) {
            return o2 - o1;
        }
    }

    /**
     * 获取添加的数的中位数
     *
     * @return 中位数 or null
     */
    public Integer getMedian() {

        // 大堆为空说明一个都没加，没有中位数
        if (maxPartQ.peek() == null) {
            return null;
        }

        // 大堆为空说明只加了一个，返回那个数就行
        if (minPartQ.peek() == null) {
            return maxPartQ.peek();
        }


        // 数目一样返回平均数，不一样返回数目多的的栈顶元素
        if (maxPartQ.size() == minPartQ.size()) {
            return (maxPartQ.peek() + minPartQ.peek()) >> 1;
        } else if (maxPartQ.size() > minPartQ.size()) {
            return maxPartQ.peek();
        } else {
            return minPartQ.peek();
        }
    }

    /**
     * 获取添加的数的中位数或者上中位数
     *
     * @return 中位数 or null
     */
    public Integer getUpMedian() {
        // 大堆为空说明一个都没加，没有中位数
        if (maxPartQ.peek() == null) {
            return null;
        }

        // 大堆为空说明只加了一个，返回那个数就行
        if (minPartQ.peek() == null) {
            return maxPartQ.peek();
        }


        // 数目一样返回上中位数，不一样返回数目多的的栈顶元素
        if (maxPartQ.size() > minPartQ.size()) {
            return maxPartQ.peek();
        } else {
            // maxPartQ.size() <= minPartQ.size()
            return minPartQ.peek();
        }
    }

    public void addElement(int e) {

        if (maxPartQ.peek() == null || maxPartQ.peek() <= e) {
            maxPartQ.add(e);
        } else {
            minPartQ.add(e);
        }

        // 判断两个堆的数目是否相差2
        if (Math.abs(maxPartQ.size() - minPartQ.size()) == 2) {
            // 谁多一个，就拿一个给另一个堆
            if (maxPartQ.size() > minPartQ.size()) {
                minPartQ.add(maxPartQ.poll());
            } else {
                maxPartQ.add(minPartQ.poll());
            }
        }
    }
}
